package mcpclient

import (
	"context"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/client/transport"
	"github.com/mark3labs/mcp-go/mcp"
)

type StdioClient struct {
	tools     []mcp.Tool        // 24 bytes (偏移 0)
	config    StdioServerConfig // 40 bytes (偏移 24)
	name      string            // 16 bytes (偏移 64)
	mu        sync.Mutex        // 8 bytes (偏移 80)
	ctx       *context.Context  // 8 bytes (偏移 88)
	client    *client.Client    // 8 bytes (偏移 96)
	taskCount atomic.Int32      // 4 bytes (偏移 104)
	open      bool              // 1 byte (偏移 108)
}

func NewStdioClient(name string, config StdioServerConfig) *StdioClient {
	stdioClient := client.NewClient(transport.NewStdio(
		config.Command, config.Env, config.Args...,
	))
	slog.Debug("NewStdioClient", "name", name)
	c := &StdioClient{
		name:      name,
		config:    config,
		client:    stdioClient,
		taskCount: atomic.Int32{},
	}
	err := c.initTools()
	if err != nil {
		slog.Error("NewStdioClient initTools", "name", name)
	}
	return c
}

func (c *StdioClient) Context(ctx *context.Context) *StdioClient {
	c.ctx = ctx // This is a placeholder for the actual context implementation
	slog.Debug("StdioClient.Context", "ctx", ctx)
	return c
}

func (c *StdioClient) connect() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.open {
		return nil
	}

	slog.Debug("StdioClient.connect", "name", c.name)

	if c.ctx == nil {
		cc := context.Background()
		c.ctx = &cc
	}

	err := c.client.Start(*c.ctx)
	if err != nil {
		slog.Error("Error starting client", "error", err)
		return err
	}

	request := mcp.InitializeRequest{}
	initResult, err := c.client.Initialize(*c.ctx, request)
	if err != nil {
		slog.Error("Error initializing client", "error", err)
		return err
	}

	if c.config.Close {
		go func() {
			defer slog.Debug("StdioClient goroutine close", "name", c.name)
			ticker := time.NewTicker(time.Second)
			defer ticker.Stop()
			for {
				if c.ctx == nil {
					slog.Info("Closing client already disconnected")
					return
				}

				select {
				case <-(*(c.ctx)).Done():
					if err := c.disconnect(); err != nil {
						slog.Error("Error disconnecting client", "error", err)
						continue
					}
					return
				case <-ticker.C:
					slog.Debug("stdioClient continue connect", "name", c.name)
				}
			}
		}()
	}

	slog.Info("Successfully initialized client", "initResult", initResult)
	c.open = true
	return nil
}

func (c *StdioClient) disconnect() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.open {
		return nil
	}

	slog.Debug("StdioClient.disconnect", "name", c.name)
	if c.taskCount.Load() > 0 {
		slog.Debug("stdioClient having task, not close", "taskCount", c.taskCount.Load())
		return nil
	}

	err := c.client.Close()
	if err != nil {
		slog.Error("Error closing client", "error", err)
		return err
	}

	c.open = false
	c.ctx = nil
	return nil
}

func (c *StdioClient) initTools() error {
	slog.Debug("StdioClient.initTools", "name", c.name)
	err := c.connect()
	if err != nil {
		return err
	}
	defer c.disconnect()

	request := mcp.ListToolsRequest{}
	toolResult, err := c.client.ListTools(*c.ctx, request)
	if err != nil {
		slog.Error("Error listing tools", "error", err)
		return err
	}
	slog.Debug("listTools", "toolResult", toolResult)
	c.tools = toolResult.Tools
	return nil
}

func (c *StdioClient) RangeTool(f func(mcp.Tool) bool) {
	for _, tool := range c.tools {
		if !f(tool) {
			return
		}
	}
}

func (c *StdioClient) CallTool(request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	slog.Debug("StdioClient.CallTool", "name", c.name, "request", request)
	defer c.disconnect()

	c.taskCount.Add(1)
	defer c.taskCount.Add(-1)
	err := c.connect()
	if err != nil {
		return nil, err
	}

	toolResult, err := c.client.CallTool(*c.ctx, request)
	if err != nil {
		slog.Error("Error calling tool", "error", err)
		return nil, err
	}

	return toolResult, nil
}
